#include "matrix3D.cuh"
#include <cstdio>
#include "timer.cuh"
#include "cuda_runtime.h"

typedef void (*SOBEL_KERNEL)(const short * __restrict__ source, short * __restrict__ dist);

__global__ void sobel_kernel_dx(const short * __restrict__ source, short * __restrict__ dist)
{
    int xthread = blockDim.x * blockIdx.x + threadIdx.x;// x coordinate
    int ythread = blockDim.y * blockIdx.y + threadIdx.y;// y coordinate
    int zthread = blockDim.z * blockIdx.z + threadIdx.z;// z coordinate

    int xcount = blockDim.x * gridDim.x;//x size
    int ycount = blockDim.y * gridDim.y;//x size
    int zcount = blockDim.z * gridDim.z;//x size
    
    int xprev = max(xthread-1, 0);
    int yprev = max(ythread-1, 0);
    int zprev = max(zthread-1, 0);

    int xnext = min(xthread+1, xcount-1);
    int ynext = min(ythread+1, ycount-1);
    int znext = min(zthread+1, zcount-1);

    unsigned int targ =  xthread + xcount * ythread + xcount * ycount * zthread;

    dist[targ] = (
        -1 * (source[xprev + xcount * yprev + xcount * ycount * zprev])     //x - 1, y - 1, z - 1
        -2 * (source[xprev + xcount * yprev + xcount * ycount * zthread])   //x - 1, y - 1, z
        -1 * (source[xprev + xcount * yprev + xcount * ycount * znext])     //x - 1, y - 1, z + 1

        -2 * (source[xprev + xcount * ythread + xcount * ycount * zprev])   //x - 1, y,     z - 1
        -4 * (source[xprev + xcount * ythread + xcount * ycount * zthread]) //x - 1, y,     z
        -2 * (source[xprev + xcount * ythread + xcount * ycount * znext])   //x - 1, y,     z + 1

        -1 * (source[xprev + xcount * ynext + xcount * ycount * zprev])     //x - 1, y + 1, z - 1
        -2 * (source[xprev + xcount * ynext + xcount * ycount * zthread])   //x - 1, y + 1, z
        -1 * (source[xprev + xcount * ynext + xcount * ycount * znext])     //x - 1, y + 1, z + 1
        
        +1 * (source[xnext + xcount * yprev + xcount * ycount * zprev])     //x + 1, y - 1, z - 1
        +2 * (source[xnext + xcount * yprev + xcount * ycount * zthread])   //x + 1, y - 1, z
        +1 * (source[xnext + xcount * yprev + xcount * ycount * znext])     //x + 1, y - 1, z + 1
        
        +2 * (source[xnext + xcount * ythread + xcount * ycount * zprev])   //x + 1, y,     z - 1
        +4 * (source[xnext + xcount * ythread + xcount * ycount * zthread]) //x + 1, y,     z
        +2 * (source[xnext + xcount * ythread + xcount * ycount * znext])   //x + 1, y,     z + 1
        
        +1 * (source[xnext + xcount * ynext + xcount * ycount * zprev])     //x + 1, y + 1, z - 1
        +2 * (source[xnext + xcount * ynext + xcount * ycount * zthread])   //x + 1, y + 1, z
        +1 * (source[xnext + xcount * ynext + xcount * ycount * znext])     //x + 1, y + 1, z + 1
        ) >> 5;
}

__global__ void sobel_kernel_dy(const short * __restrict__ source, short * __restrict__ dist)
{
    int xthread = blockDim.x * blockIdx.x + threadIdx.x;// x coordinate
    int ythread = blockDim.y * blockIdx.y + threadIdx.y;// y coordinate
    int zthread = blockDim.z * blockIdx.z + threadIdx.z;// z coordinate

    int xcount = blockDim.x * gridDim.x;//x size
    int ycount = blockDim.y * gridDim.y;//x size
    int zcount = blockDim.z * gridDim.z;//x size
    
    int xprev = max(xthread-1, 0);
    int yprev = max(ythread-1, 0);
    int zprev = max(zthread-1, 0);

    int xnext = min(xthread+1, xcount-1);
    int ynext = min(ythread+1, ycount-1);
    int znext = min(zthread+1, zcount-1);

    unsigned int targ =  xthread + xcount * ythread + xcount * ycount * zthread;

    dist[targ] = (
        -1 * (source[xprev + xcount * yprev + xcount * ycount * zprev])     //x - 1, y - 1, z - 1
        -2 * (source[xprev + xcount * yprev + xcount * ycount * zthread])   //x - 1, y - 1, z
        -1 * (source[xprev + xcount * yprev + xcount * ycount * znext])     //x - 1, y - 1, z + 1

        +1 * (source[xprev + xcount * ynext + xcount * ycount * zprev])     //x - 1, y + 1, z - 1
        +2 * (source[xprev + xcount * ynext + xcount * ycount * zthread])   //x - 1, y + 1, z
        +1 * (source[xprev + xcount * ynext + xcount * ycount * znext])     //x - 1, y + 1, z + 1

        -2 * (source[xthread + xcount * yprev + xcount * ycount * zprev])   //x,     y,     z - 1
        -4 * (source[xthread + xcount * yprev + xcount * ycount * zthread]) //x,     y,     z
        -2 * (source[xthread + xcount * yprev + xcount * ycount * znext])   //x,     y,     z + 1

        +2 * (source[xthread + xcount * ynext + xcount * ycount * zprev])   //x,     y + 1, z - 1
        +4 * (source[xthread + xcount * ynext + xcount * ycount * zthread]) //x,     y + 1, z
        +2 * (source[xthread + xcount * ynext + xcount * ycount * znext])   //x,     y + 1, z + 1
        
        -1 * (source[xnext + xcount * yprev + xcount * ycount * zprev])     //x + 1, y - 1, z - 1
        -2 * (source[xnext + xcount * yprev + xcount * ycount * zthread])   //x + 1, y - 1, z
        -1 * (source[xnext + xcount * yprev + xcount * ycount * znext])     //x + 1, y - 1, z + 1
        
        +1 * (source[xnext + xcount * ynext + xcount * ycount * zprev])     //x + 1, y + 1, z - 1
        +2 * (source[xnext + xcount * ynext + xcount * ycount * zthread])   //x + 1, y + 1, z
        +1 * (source[xnext + xcount * ynext + xcount * ycount * znext])     //x + 1, y + 1, z + 1
        ) >> 5;
}

__global__ void sobel_kernel_dz(const short * __restrict__ source, short * __restrict__ dist)
{
    int xthread = blockDim.x * blockIdx.x + threadIdx.x;// x coordinate
    int ythread = blockDim.y * blockIdx.y + threadIdx.y;// y coordinate
    int zthread = blockDim.z * blockIdx.z + threadIdx.z;// z coordinate

    int xcount = blockDim.x * gridDim.x;//x size
    int ycount = blockDim.y * gridDim.y;//x size
    int zcount = blockDim.z * gridDim.z;//x size
    
    int xprev = max(xthread-1, 0);
    int yprev = max(ythread-1, 0);
    int zprev = max(zthread-1, 0);

    int xnext = min(xthread+1, xcount-1);
    int ynext = min(ythread+1, ycount-1);
    int znext = min(zthread+1, zcount-1);

    unsigned int targ =  xthread + xcount * ythread + xcount * ycount * zthread;

    dist[targ] = (
        -1 * (source[xprev + xcount * yprev + xcount * ycount * zprev])     //x - 1, y - 1, z - 1
        -2 * (source[xprev + xcount * ythread + xcount * ycount * zprev])   //x - 1, y,     z - 1
        -1 * (source[xprev + xcount * ynext + xcount * ycount * zprev])     //x - 1, y + 1, z - 1

        +1 * (source[xprev + xcount * yprev + xcount * ycount * znext])     //x - 1, y -1,  z + 1
        +2 * (source[xprev + xcount * ythread + xcount * ycount * znext])   //x - 1, y,     z + 1
        +1 * (source[xprev + xcount * ynext + xcount * ycount * znext])     //x - 1, y +1,  z + 1

        -2 * (source[xthread + xcount * yprev + xcount * ycount * zprev])   //x,     y - 1, z - 1
        -4 * (source[xthread + xcount * ythread + xcount * ycount * zprev]) //x,     y,     z - 1
        -2 * (source[xthread + xcount * ynext + xcount * ycount * zprev])   //x,     y + 1, z - 1
        
        +2 * (source[xthread + xcount * yprev + xcount * ycount * znext])   //x,     y - 1, z + 1
        +4 * (source[xthread + xcount * ythread + xcount * ycount * znext]) //x,     y,     z + 1
        +2 * (source[xthread + xcount * ynext + xcount * ycount * znext])   //x,     y + 1, z + 1
        
        -1 * (source[xnext + xcount * yprev + xcount * ycount * zprev])     //x + 1, y - 1, z - 1
        -2 * (source[xnext + xcount * ythread + xcount * ycount * zprev])   //x + 1, y,     z - 1
        -1 * (source[xnext + xcount * ynext + xcount * ycount * zprev])     //x + 1, y + 1, z - 1
        
        +1 * (source[xnext + xcount * yprev + xcount * ycount * znext])     //x + 1, y - 1, z + 1
        +2 * (source[xnext + xcount * ythread + xcount * ycount * znext])   //x + 1, y,     z + 1
        +1 * (source[xnext + xcount * ynext + xcount * ycount * znext])     //x + 1, y + 1, z + 1
        ) >> 5;
}

void cuda_sobel( const short *source, short *dist, int xcount, int ycount, int zcount, int type)
{
    short* d_source;
    short* d_dist;
    int model_size = xcount * ycount * zcount;
    cudaError_t cudaStatus;

    cudaStatus = cudaMalloc((void**)&d_source, model_size * sizeof(short));
    if (cudaStatus != cudaSuccess)
    {
        printf("cudaMalloc failed for model memory request!\n");
        return;
    }

    cudaStatus = cudaMalloc((void**)&d_dist, model_size * sizeof(short));
    if (cudaStatus != cudaSuccess)
    {
        printf("cudaMalloc failed for histogram memory request!\n");
        cudaFree(d_source);
        return;
    }

    cudaStatus = cudaMemcpy(d_source, source, model_size * sizeof(short), cudaMemcpyHostToDevice);
    if (cudaStatus != cudaSuccess)
    {
        printf("cudaMemcpy failed for sending model data on GPU!");
        cudaFree(d_source);
        cudaFree(d_dist);
    }

    dim3 threads(32, 4, 8);
    dim3 grid = dim3(xcount >> 5, ycount >> 2, zcount >> 3);
    SOBEL_KERNEL kernel;
    switch(type)
    {
        case 0:{kernel = sobel_kernel_dx; break;}
        case 1:{kernel = sobel_kernel_dy; break;}
        case 2:{kernel = sobel_kernel_dz; break;}
    }

    Timer t;
    cudaEvent_t start = t.go();

    kernel<<<grid, threads>>>(d_source, d_dist);

    cudaStatus = cudaDeviceSynchronize();

    printf("Sobel calculation complete after : %f Ms \n", t.get(start));

    if (cudaStatus != cudaSuccess)
    {
        printf("cudaDeviceSynchronize returned error code %d after launching addKernel!\n", cudaStatus);
        goto Error;
    }

    cudaStatus = cudaMemcpy(dist, d_dist, model_size * sizeof(short), cudaMemcpyDeviceToHost);
    if (cudaStatus != cudaSuccess)
    {
        printf("cudaMemcpy failed!");
        goto Error;
    }

Error:
    cudaFree(d_source);
    cudaFree(d_dist);
}

